import os
import random

# 设置随机种子以保证结果的可重复性
random.seed(42)

# 直接设置图像数据集路径
image_dir = 'VOC2028_short/images1'  # 替换成自己的数据集路径

# 获取所有图片的文件名
all_images = [os.path.join(image_dir, img) for img in os.listdir(image_dir)
              if img.lower().endswith(('png', 'jpg', 'jpeg', 'bmp'))]
"""以上这句复合代码的解释，二选一背
# 获取图像数据集目录下的所有文件名
all_file_names = os.listdir(image_dir)

# 使用循环方式
all_images = []
for img in all_file_names:  # 对于“所有文件名”中的每个图像文件
    if img.lower().endswith(('png', 'jpg', 'jpeg', 'bmp')): 
        all_images.append(os.path.join(image_dir, img))  # 获得图片路径，加入到数组中
"""

# 打乱图片列表
random.shuffle(all_images)

# 划分数据集
train_split = 0.7  # 指定训练集百分比
val_split = 0.15  # 指定验证集百分比
train_count = int(len(all_images) * train_split)  # 训练集实际长度
val_count = int(len(all_images) * val_split)  # 验证集实际长度

train_images = all_images[:train_count]  # 训练集包含的图片路径
val_images = all_images[train_count:train_count + val_count]  # 验证集包含的图片路径
test_images = all_images[train_count + val_count:]  # 剩余为测试集


# 保存到txt文件
def save_to_file(images, filename):
    with open(filename, 'w') as f:
        f.write('\n'.join(images) + '\n')


# 获取父文件夹路径
parent_dir = os.path.dirname(image_dir)

# 保存文件
save_to_file(train_images, os.path.join(parent_dir, 'train.txt'))
save_to_file(val_images, os.path.join(parent_dir, 'val.txt'))
save_to_file(test_images, os.path.join(parent_dir, 'test.txt'))
save_to_file(train_images + val_images, os.path.join(parent_dir, 'trainval.txt'))

print(f"数据集划分完成: {len(all_images)} 张图片")
print(f"训练集数量: {len(train_images)}")
print(f"验证集数量: {len(val_images)}")
print(f"测试集数量: {len(test_images)}")
print(f"训练集+验证集数量: {len(train_images + val_images)}")

# 创建data.yaml文件
print("开始生成配置文件...")
classes = 2  # 根据实际需要修改类别数量 # 根据实际类别名称修改
data_yaml = f"""
train: {os.path.abspath(os.path.join(parent_dir, 'train.txt'))}
val: {os.path.abspath(os.path.join(parent_dir, 'val.txt'))}
test: {os.path.abspath(os.path.join(parent_dir, 'test.txt'))}

nc: {classes}
names: ['class1', 'class2'] 
"""
# names 根据实际类别名称修改

with open(os.path.join(parent_dir, 'data.yaml'), 'w') as f:
    f.write(data_yaml)
